// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/vector/decoded_vector.h>
#include <pollux/vector/simple_vector.h>

namespace kumo::pollux {
    // A set of unique, increasing row numbers for use as qualifying
    // set. This is logically interchangeable with
    // SelectivityVector. Since column reading does not have frequent bit
    // operations but always loops over often sparse positions, an array
    // of positions is more convenient. melon::Range is also faster to
    // pass and access than the indirections and smart pointers that may
    // be involved in SelectivityVector.
    using RowSet = melon::Range<const vector_size_t *>;

    // Defines a per-value callback to apply to values when loading a
    // LazyVector. This enables pushing down an arbitrary operation into
    // e.g. table scan when loading columns of scalar data types. In this
    // way one can bypass copying data into a vector before use.
    class ValueHook {
    public:
        // Type and constants for identifying specific hooks.  Loaders may have
        // hook-specialized template instantiations for some operations.
        enum Kind {
            kGeneric,
            kBigintSum,
            kBigintSumOverflow,
            kDoubleSum,
            kBigintMax,
            kBigintMin,
            kFloatingPointMax,
            kFloatingPointMin,
        };

        static constexpr bool kSkipNulls = true;

        virtual ~ValueHook() = default;

        virtual bool acceptsNulls() const {
            return false;
        }

        virtual Kind kind() const {
            return kGeneric;
        }

        virtual void addNull(vector_size_t /*index*/) {
        }

        virtual void addValue(vector_size_t /*row*/, int64_t /*value*/) {
            POLLUX_UNSUPPORTED();
        }

        virtual void addValue(vector_size_t /*row*/, int128_t /*value*/) {
            POLLUX_UNSUPPORTED();
        }

        virtual void addValue(vector_size_t /*row*/, float /*value*/) {
            POLLUX_UNSUPPORTED();
        }

        virtual void addValue(vector_size_t /*row*/, double /*value*/) {
            POLLUX_UNSUPPORTED();
        }

        virtual void addValue(vector_size_t /*row*/, melon::StringPiece /*value*/) {
            POLLUX_UNSUPPORTED();
        }

        // Fallback implementation of bulk path for addValues.  Actual hooks are
        // expected to override these if they are not inlined in reader.
        virtual void
        addValues(const vector_size_t *rows, const bool *values, vector_size_t size) {
            for (auto i = 0; i < size; ++i) {
                addValueTyped(rows[i], values[i]);
            }
        }

        virtual void addValues(
            const vector_size_t *rows,
            const int8_t *values,
            vector_size_t size) {
            for (auto i = 0; i < size; ++i) {
                addValueTyped(rows[i], values[i]);
            }
        }

        virtual void addValues(
            const vector_size_t *rows,
            const int16_t *values,
            vector_size_t size) {
            for (auto i = 0; i < size; ++i) {
                addValueTyped(rows[i], values[i]);
            }
        }

        virtual void addValues(
            const vector_size_t *rows,
            const int32_t *values,
            vector_size_t size) {
            for (auto i = 0; i < size; ++i) {
                addValueTyped(rows[i], values[i]);
            }
        }

        virtual void addValues(
            const vector_size_t *rows,
            const int64_t *values,
            vector_size_t size) {
            for (auto i = 0; i < size; ++i) {
                addValue(rows[i], values[i]);
            }
        }

        virtual void addValues(
            const vector_size_t *rows,
            const int128_t *values,
            vector_size_t size) {
            for (auto i = 0; i < size; ++i) {
                addValue(rows[i], values[i]);
            }
        }

        virtual void addValues(
            const vector_size_t *rows,
            const float *values,
            vector_size_t size) {
            for (auto i = 0; i < size; ++i) {
                addValue(rows[i], values[i]);
            }
        }

        virtual void addValues(
            const vector_size_t *rows,
            const double *values,
            vector_size_t size) {
            for (auto i = 0; i < size; ++i) {
                addValue(rows[i], values[i]);
            }
        }

        virtual void addValues(
            const vector_size_t *rows,
            const StringView *values,
            vector_size_t size) {
            for (auto i = 0; i < size; ++i) {
                addValue(rows[i], values[i]);
            }
        }

        template<typename T>
        void addValueTyped(vector_size_t row, T value) {
            if constexpr (std::is_integral_v<T> && sizeof(T) < sizeof(int64_t)) {
                addValue(row, static_cast<int64_t>(value));
            } else {
                addValue(row, value);
            }
        }
    };

    // Produces values for a LazyVector for a set of positions.
    class VectorLoader {
    public:
        virtual ~VectorLoader() = default;

        // Produces the lazy values for 'rows' and if 'hook' is non-nullptr,
        // calls hook on each. If 'hook' is nullptr, sets '*result' to a
        // vector that contains the values for 'rows'. 'rows' must be a
        // subset of the rows that were intended to be loadable when the
        // loader was created. This may be called once in the lifetime of
        // 'this'.
        // Notes: Implementations of this class should ensure:
        // 1.‘result’ is unique before mutating it.
        // 2. result’ size is at least resultSize.
        void load(
            RowSet rows,
            ValueHook *hook,
            vector_size_t resultSize,
            VectorPtr *result);

        // Converts 'rows' into a RowSet and calls load(). Provided for
        // convenience in loading LazyVectors in expression evaluation.
        void load(
            const SelectivityVector &rows,
            ValueHook *hook,
            vector_size_t resultSize,
            VectorPtr *result);

    protected:
        virtual void loadInternal(
            RowSet rows,
            ValueHook *hook,
            vector_size_t resultSize,
            VectorPtr *result) = 0;
    };

    // Vector class which produces values on first use. This is used for
    // loading columns on demand. This allows eliding load of
    // columns which have all values filtered out by e.g. joins or which
    // are never referenced due to conditionals in projection. If the call
    // site known that only a subset of the positions in the vector will
    // ever be accessed, loading can be limited to these positions. This
    // also allows pushing down computation into loading a column, hence
    // bypassing materialization into a vector.
    // Unloaded LazyVectors should be referenced only by one top-level vector.
    // Otherwise, it runs the risk of being loaded for different set of rows by each
    // top-level vector.
    class LazyVector : public BaseVector {
    public:
        static constexpr const char *kCpuNanos = "dataSourceLazyCpuNanos";
        static constexpr const char *kWallNanos = "dataSourceLazyWallNanos";
        static constexpr const char *kInputBytes = "dataSourceLazyInputBytes";

        LazyVector(
            pollux::memory::MemoryPool *pool,
            TypePtr type,
            vector_size_t size,
            std::unique_ptr<VectorLoader> &&loader,
            VectorPtr &&vector = nullptr)
            : BaseVector(
                  pool,
                  std::move(type),
                  VectorEncoding::Simple::LAZY,
                  BufferPtr(nullptr),
                  size),
              loader_(std::move(loader)),
              vector_(std::move(vector)) {
        }

        void reset(std::unique_ptr<VectorLoader> &&loader, vector_size_t size) {
            BaseVector::length_ = size;
            loader_ = std::move(loader);
            allLoaded_ = false;
            containsLazyAndIsWrapped_ = false;
            reset_nulls();
        }

        inline bool isLoaded() const {
            return allLoaded_;
        }

        // Loads the positions in 'rows' into loadedVector_. If 'hook' is
        // non-nullptr, the hook is instead called on the values and
        // loaded_vector is not updated. This method is const because call
        // sites often have a const VaseVector. Lazy construction is
        // logically not a mutation.
        void load(RowSet rows, ValueHook *hook) const;

        std::optional<int32_t> compare(
            const BaseVector *other,
            vector_size_t index,
            vector_size_t otherIndex,
            CompareFlags flags) const override {
            return loaded_vector()->compare(other, index, otherIndex, flags);
        }

        uint64_t hash_value_at(vector_size_t index) const override {
            return loaded_vector()->hash_value_at(index);
        }

        std::unique_ptr<SimpleVector<uint64_t> > hash_all() const override {
            return loaded_vector()->hash_all();
        }

        const BaseVector *loaded_vector() const override {
            return loaded_vector_shared().get();
        }

        BaseVector *loaded_vector() override {
            return loaded_vector_shared().get();
        }

        // Returns a shared_ptr to the vector holding the values. If vector is not
        // loaded, loads all the rows, otherwise returns the loaded vector which can
        // have partially loaded rows.
        VectorPtr &loaded_vector_shared() {
            loadVectorInternal();
            return vector_;
        }

        const VectorPtr &loaded_vector_shared() const {
            loadVectorInternal();
            return vector_;
        }

        const BaseVector *wrapped_vector() const override {
            return loaded_vector()->wrapped_vector();
        }

        vector_size_t wrapped_index(vector_size_t index) const override {
            return loaded_vector()->wrapped_index(index);
        }

        const BufferPtr &wrap_info() const override {
            return loaded_vector()->wrap_info();
        }

        bool is_scalar() const override {
            return type()->isPrimitiveType() || type()->is_opaque();
        }

        bool may_have_nulls() const override {
            return loaded_vector()->may_have_nulls();
        }

        bool may_have_nulls_recursive() const override {
            return loaded_vector()->may_have_nulls_recursive();
        }

        bool is_null_at(vector_size_t index) const override {
            return loaded_vector()->is_null_at(index);
        }

        bool contains_null_at(vector_size_t index) const override {
            return loaded_vector()->contains_null_at(index);
        }

        uint64_t retained_size() const override {
            return isLoaded()
                       ? loaded_vector()->retained_size()
                       : BaseVector::retained_size();
        }

        /// Returns zero if vector has not been loaded yet.
        uint64_t estimate_flat_size() const override {
            return isLoaded() ? loaded_vector()->estimate_flat_size() : 0;
        }

        std::string toString(vector_size_t index) const override {
            return loaded_vector()->toString(index);
        }

        VectorPtr slice(vector_size_t offset, vector_size_t length) const override;

        // Loads 'rows' of 'vector'. 'vector' may be an arbitrary wrapping
        // of a LazyVector. 'rows' are translated through the wrappers. If
        // there is no LazyVector inside 'vector', this has no
        // effect.
        static void ensureLoadedRows(
            const VectorPtr &vector,
            const SelectivityVector &rows);

        // as ensureLoadedRows, above, but takes a scratch DecodedVector and
        // SelectivityVector as arguments to enable reuse.
        static void ensureLoadedRows(
            const VectorPtr &vector,
            const SelectivityVector &rows,
            DecodedVector &decoded,
            SelectivityVector &baseRows);

        void validate(const VectorValidateOptions &options) const override;

        VectorPtr copy_preserve_encodings(
            pollux::memory::MemoryPool *pool = nullptr) const override {
            POLLUX_CHECK(isLoaded());
            return loaded_vector()->copy_preserve_encodings(pool);
        }

    private:
        static void ensureLoadedRowsImpl(
            const VectorPtr &vector,
            DecodedVector &decoded,
            const SelectivityVector &rows,
            SelectivityVector &baseRows);

        void loadVectorInternal() const;

        std::unique_ptr<VectorLoader> loader_;

        // True if all values are loaded.
        mutable tsan_atomic<bool> allLoaded_{false};
        // Vector to hold loaded values. This may be present before load for
        // reuse. If loading is with ValueHook, this will not be created.
        mutable VectorPtr vector_;
    };
} // namespace kumo::pollux
